package com.keguoyu.easymvp.ioc.origin

import com.google.common.base.Preconditions
import com.google.common.io.ByteStreams
import com.google.common.io.Closeables
import com.keguoyu.easymvp.ioc.internal.Invocation
import com.keguoyu.easymvp.ioc.internal.InvokerInfos
import com.keguoyu.easymvp.ioc.internal.util.toClassName
import com.keguoyu.easymvp.ioc.internal.util.toClassPath
import javassist.ClassPool
import javassist.CtClass
import javassist.CtMethod
import org.apache.commons.io.FileUtils
import java.io.ByteArrayInputStream
import java.io.File
import java.io.FileOutputStream
import java.io.InputStream
import java.util.jar.JarFile
import java.util.jar.JarOutputStream
import java.util.zip.ZipEntry

class FileModifier(private val mInfos: InvokerInfos,
                   private val mIsDebug: Boolean,
                   files: FileSet) {
  private val mPool = ClassPool(true)

  init {
    files.loopFiles {
      mPool.appendClassPath(it.absolutePath)
    }
    files.loopJars {
      mPool.appendClassPath(it.path)
    }
  }

  fun modify(files: Collection<File>, jars: Collection<File>) {
    files.modifyClass({ mInfos.classInFile(it.absolutePath) }, { dest, classes -> writeFile(dest, classes) })
    jars.modifyClass({ mInfos.classInJar(it.absolutePath) }, { dest, classes -> writeJar(dest, classes) })
  }

  private fun writeJar(file: File, classes: Map<String, ByteArray>) {
    if (classes.isEmpty()) {
      return
    }
    val tmpFile = File(file.getParent(), file.getName() + ".tmp")
    if (tmpFile.exists()) {
      tmpFile.delete()
    }
    val output = JarOutputStream(FileOutputStream(tmpFile))
    val jar = JarFile(file)
    jar.entries().iterator().forEach { entry ->
      var input: InputStream? = null
      if (entry.getName().endsWith(".class")) {
        val className = entry.toClassName()
        val bytes = classes[className]
        if (bytes != null) {
          input = ByteArrayInputStream(bytes)
          if (mIsDebug) {
            val fout = File(file.getParent(), "${className}.modified.class")
            if (fout.exists()) {
              fout.delete()
              fout.createNewFile()
            }
            ByteStreams.copy(ByteArrayInputStream(bytes), FileOutputStream(fout))
          }
        }
      }
      if (input == null) {
        input = jar.getInputStream(entry)
      }
      val outEntry = ZipEntry(entry.getName())
      output.putNextEntry(outEntry)
      ByteStreams.copy(input!!, output)
      Closeables.close(input, true)
      output.closeEntry()
    }
    Closeables.close(output, true)

    if (mIsDebug) {
      file.renameTo(File(file.getParent(), "${file.getName()}.backup.jar"))
    }
    if (file.exists()) {
      file.delete()
    }
    tmpFile.renameTo(file)
  }

  private fun writeFile(fout: File, classes: Map<String, ByteArray>) {
    if (classes.isEmpty()) {
      return
    }
    if (fout.exists()) {
      fout.delete()
    }
    fout.createNewFile()
    Preconditions.checkState(classes.size == 1, "${classes.size} 一个 class 文件只能有一个类 ${fout.absolutePath}")
    val (clazz, bytes) = classes.entries.elementAt(0)
    Preconditions.checkState(fout.absolutePath.endsWith(clazz.toClassPath()))
    FileUtils.writeByteArrayToFile(fout, bytes)
  }

  private fun buildClassBytes(clazz: String, invocations: Collection<Invocation>): ByteArray {
    // 增加排序, 保证每次编译生成的结果类是一样的
    val sortedInvocations = invocations.sortedWith(Comparator<Invocation>() { a, b ->
      a.target.className.compareTo(b.target.className)
    })
    val invoker = mPool.get(clazz)
    if (invoker.isFrozen()) {
      invoker.defrost()
    }
    sortedInvocations.forEach { invocation ->
      val targetClass = invocation.target.className
      mPool.importPackage(targetClass)
      val types = buildParamTypes(invocation)

      var invokerMethod: CtMethod
      if (types.isEmpty()) {
        invokerMethod = invoker.getDeclaredMethod(invocation.invoker.methodName)
      } else {
        invokerMethod = invoker.getDeclaredMethod(invocation.invoker.methodName, types)
      }

      invokerMethod.insertAfter(buildSrc(invocation))
    }
    val bytes = invoker.toBytecode()
    invoker.detach()
    return bytes
  }

  private fun buildParamTypes(invocation: Invocation): Array<CtClass> {
    val types = Array<CtClass>(invocation.params.size) { _ -> CtClass.voidType }
    invocation.params.forEachIndexed { index, type ->
      // java 原生类型不在原文件列表中
      mPool.importPackage(type)
      types.set(index, mPool.get(type))
    }
    return types
  }

  private fun buildSrc(invocation: Invocation): String {
    var src = "${invocation.target.className}.${invocation.target.methodName}("
    if (!invocation.isStatic) {
      src += "\$0"
    }
    if (!invocation.params.isEmpty()) {
      if (src.contains("\$")) {
        src += ","
      }
      src += "\$\$"
    }
    return src + ");"
  }

  private fun Collection<File>.modifyClass(fetcher: (File) -> Collection<String>,
                                           writer: (File, Map<String, ByteArray>) -> Unit) {
    forEach {
      val classes = fetcher(it)
      val modifiedClasses = HashMap<String, ByteArray>()
      classes.forEach loop@{ clazz ->
        if (mInfos.invocationForClass(clazz)?.isEmpty() ?: true) {
          return@loop
        }
        modifiedClasses[clazz] = buildClassBytes(clazz, mInfos.invocationForClass(clazz)!!)
      }
      if (modifiedClasses.isEmpty()) {
        return@forEach
      }
      writer(it, modifiedClasses)
    }
  }
}